import asyncio
from langchain_core.documents import Document

def test_summarize_entity_1(cp):
    keyword, descriptions, semantical_merged, summary = asyncio.run(cp._summarize({'id':'市领导', 'degree': 1000}))

    assert '市领导'==keyword

def test_summarize_entity_2(cp):
    keyword, descriptions, semantical_merged, summary = asyncio.run(cp._summarize({'id':'陈金虎', 'degree': 1000}))

    assert '陈金虎'==keyword

def test_summarize_entity_3(cp):
    keyword, descriptions, semantical_merged, summary = asyncio.run(cp._summarize({'id':'现代传媒中心', 'degree': 1000}))
    assert '现代传媒中心'==keyword

def test_extract_ner(cp):
    documents = [
{
    'metadata':{
       'reported_on': "2025/01/01",
       "chunk_id": "01",
       "title": "陈金虎周伟走访财税金融机构",
       "cameraman": ["朱佳明"],
       "reporter": ["刘晔"],
       "voiceover": ["赖岳"]
    },
    "page_content":"""
持续提升金融服务质效 激发市场主体活力竞争力
 【口播】 
    昨天（12月31日），市委书记陈金虎、代市长周伟带队走访市财税部门和部分金融机构，走访慰问一线工作人员，强调要深入学习贯彻习近平总书记对江苏工作重要讲话精神，认真落实中央经济工作会议和省委全会、省委经济工作会议精神，按照市委十三届十次全会部署，深化财税、金融体制改革，创新金融服务，服务实体经济，高质量完成“十四五”规划目标任务，为实现“十五五”良好开局打牢基础。
 【正文】 
    陈金虎、周伟先后前往农业银行常州分行、招商银行常州分行，走进服务大厅，详细了解机构运行、金融服务等工作开展情况，对各金融机构对城市经济社会发展的关心支持表示感谢。市领导希望各金融机构进一步发挥金融“活水”作用，精准滴灌千企万户所需，助力新质生产力发展。要完善服务体系、创新金融产品，提升金融服务的精准性和有效性，更好服务实体经济、助力重大项目、推进乡村振兴，以金融高质量发展助力促消费扩内需。要大力弘扬劳模精神、工匠精神，打造更多知识型、专业型、技能型、创新型的金融职工队伍。
    在市税务局，陈金虎、周伟与工作人员亲切交谈，向大家致以诚挚问候和新年祝福。市领导详细了解年终工作开展情况，强调要充分发挥税收职能作用，积极服务市场主体，推动税收服务更加规范化、人性化，为全市经济社会高质量发展作出贡献。
    在认真听取全市金融、财政、税务运行情况汇报后，市领导指出，金融是国民经济的血脉。全市财税金融系统要聚焦主责主业，坚定信心决心，持续提升金融服务质效，引导更多金融“活水”赋能实体经济，激发市场主体活力、竞争力，政企同频共振、同向发力。要开好“新春第一会”全市企业发展大会，出台专题政策文件和支持举措，持续优化营商环境，提振企业信心，营造良好发展氛围。
    市领导李林、杭勇、徐军和市政府秘书长周承涛参加走访。
"""
}
    ]
    results = asyncio.run(cp._extract_ner(documents))
    assert 1 == len(results)
    descriptions = results[0]["descriptions"]
    assert len(descriptions) > 3

def test_ollama():
    from openai import OpenAI

    client = OpenAI(
        base_url = 'http://localhost:11434/v1',
        api_key='ollama', # required, but unused
    )

    response = client.chat.completions.create(
    model="15b",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "用JSON格式输出中国历史朝代名称和存亡年代，{'name':'明朝', start='公元1368年', end='公元1644年'}"},
    ]
    )
    print(response.choices[0].message.content)